# Setup environment ----
library(parallel)
library(data.table)
library(glmnet)
RNGkind("L'Ecuyer-CMRG")
source("R/functions.R")
results_path <- "results/replication/"
load(paste0(results_path, "house_analysis_data.RData"))
load(paste0(results_path, "senate_analysis_data.RData"))

# House LASSO ----
house_lasso_model <- glmnet(x = hou_X, y = hou_y,
  lambda = cv.glmnet(x = hou_X, y = hou_y)$lambda.1se)
cv_lasso_hou <- cv_lasso(hou_folds, hou_X, hou_y)
house_lasso_data <- data.table(
  dem_spend_adv = hou_X[, "dem_spend_adv"],
  d_dem_spend_adv = ((
    predict(house_lasso_model, Xp1H) -
      predict(house_lasso_model, Xm1H)) /
      (2 * dxH))[, 1],
  incl = hou_X[, "bottom_tail"] + hou_X[, "top_tail"] == 0)
save(house_lasso_model, cv_lasso_hou, house_lasso_data,
  file = paste0(results_path, "house_lasso.RData"))

# Senate LASSO ----
senate_lasso_model <- glmnet(x = sen_X, y = sen_y,
  lambda = cv.glmnet(x = sen_X, y = sen_y)$lambda.1se)
cv_lasso_sen <- cv_lasso(sen_folds, sen_X, sen_y)
senate_lasso_data <- data.table(
  dem_spend_adv = sen_X[, "dem_spend_adv"],
  d_dem_spend_adv = ((
    predict(senate_lasso_model, Xp1S) -
      predict(senate_lasso_model, Xm1S)) /
      (2 * dxS))[, 1],
  incl = sen_X[, "bottom_tail"] + sen_X[, "top_tail"] == 0)
save(senate_lasso_model, cv_lasso_sen, senate_lasso_data,
  file = paste0(results_path, "senate_lasso.RData"))

# House LASSO with polynominal ----
hou_X_poly <- cbind(hou_X,
  hou_X[, "dem_spend_adv"]^2, hou_X[, "dem_spend_adv"]^3,
  hou_X[, "dem_spend_adv"]^4,
  hou_X[, "bottom_tail_DSA"]^2, hou_X[, "bottom_tail_DSA"]^3,
  hou_X[, "bottom_tail_DSA"]^4,
  hou_X[, "top_tail_DSA"]^2, hou_X[, "top_tail_DSA"]^3,
  hou_X[, "top_tail_DSA"]^4)
house_lasso_poly_model <- glmnet(x = hou_X_poly, y = hou_y,
  lambda = cv.glmnet(x = hou_X_poly, y = hou_y)$lambda.1se)
cv_lasso_poly_hou <- cv_lasso(hou_folds, hou_X_poly, hou_y)
Xp1H_poly <- cbind(Xp1H,
  Xp1H[, "dem_spend_adv"]^2, Xp1H[, "dem_spend_adv"]^3,
  Xp1H[, "dem_spend_adv"]^4,
  Xp1H[, "bottom_tail_DSA"]^2, Xp1H[, "bottom_tail_DSA"]^3,
  Xp1H[, "bottom_tail_DSA"]^4,
  Xp1H[, "top_tail_DSA"]^2, Xp1H[, "top_tail_DSA"]^3,
  Xp1H[, "top_tail_DSA"]^4)
Xm1H_poly <- cbind(Xm1H,
  Xm1H[, "dem_spend_adv"]^2, Xm1H[, "dem_spend_adv"]^3,
  Xm1H[, "dem_spend_adv"]^4,
  Xm1H[, "bottom_tail_DSA"]^2, Xm1H[, "bottom_tail_DSA"]^3,
  Xm1H[, "bottom_tail_DSA"]^4,
  Xm1H[, "top_tail_DSA"]^2, Xm1H[, "top_tail_DSA"]^3,
  Xm1H[, "top_tail_DSA"]^4)
house_lasso_poly_data <- data.table(
  dem_spend_adv = hou_X_poly[, "dem_spend_adv"],
  d_dem_spend_adv = ((
    predict(house_lasso_poly_model, Xp1H_poly) -
      predict(house_lasso_poly_model, Xm1H_poly)) /
      (2 * dxH))[, 1],
  incl = hou_X_poly[, "bottom_tail"] + hou_X_poly[, "top_tail"] == 0)
save(house_lasso_poly_model, cv_lasso_poly_hou, house_lasso_poly_data,
  file = paste0(results_path, "house_lasso_poly.RData"))
# Senate LASSO with polynominal ----
sen_X_poly <- cbind(sen_X,
  sen_X[, "dem_spend_adv"]^2, sen_X[, "dem_spend_adv"]^3,
  sen_X[, "dem_spend_adv"]^4,
  sen_X[, "bottom_tail_DSA"]^2, sen_X[, "bottom_tail_DSA"]^3,
  sen_X[, "bottom_tail_DSA"]^4,
  sen_X[, "top_tail_DSA"]^2, sen_X[, "top_tail_DSA"]^3,
  sen_X[, "top_tail_DSA"]^4)
senate_lasso_poly_model <- glmnet(x = sen_X_poly, y = sen_y,
  lambda = cv.glmnet(x = sen_X_poly, y = sen_y)$lambda.1se)
cv_lasso_poly_sen <- cv_lasso(sen_folds, sen_X_poly, sen_y)
Xp1S_poly <- cbind(Xp1S,
  Xp1S[, "dem_spend_adv"]^2, Xp1S[, "dem_spend_adv"]^3,
  Xp1S[, "dem_spend_adv"]^4,
  Xp1S[, "bottom_tail_DSA"]^2, Xp1S[, "bottom_tail_DSA"]^3,
  Xp1S[, "bottom_tail_DSA"]^4,
  Xp1S[, "top_tail_DSA"]^2, Xp1S[, "top_tail_DSA"]^3,
  Xp1S[, "top_tail_DSA"]^4)
Xm1S_poly <- cbind(Xm1S,
  Xm1S[, "dem_spend_adv"]^2, Xm1S[, "dem_spend_adv"]^3,
  Xm1S[, "dem_spend_adv"]^4,
  Xm1S[, "bottom_tail_DSA"]^2, Xm1S[, "bottom_tail_DSA"]^3,
  Xm1S[, "bottom_tail_DSA"]^4,
  Xm1S[, "top_tail_DSA"]^2, Xm1S[, "top_tail_DSA"]^3,
  Xm1S[, "top_tail_DSA"]^4)
senate_lasso_poly_data <- data.table(
  dem_spend_adv = sen_X_poly[, "dem_spend_adv"],
  d_dem_spend_adv = ((
    predict(senate_lasso_poly_model, Xp1S_poly) -
      predict(senate_lasso_poly_model, Xm1S_poly)) /
      (2 * dxS))[, 1],
  incl = sen_X_poly[, "bottom_tail"] + sen_X_poly[, "top_tail"] == 0)
save(senate_lasso_poly_model, cv_lasso_poly_sen, senate_lasso_poly_data,
  file = paste0(results_path, "sen_lasso_poly.RData"))
